#!/usr/bin/env python
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
"""
Detection Training Script.

This scripts reads a given config file and runs the training or evaluation.
It is an entry point that is made to train standard models in detectron2.

In order to let one script support training of many models,
this script contains logic that are specific to these built-in models and therefore
may not be suitable for your own project.
For example, your research project perhaps only needs a single "evaluator".

Therefore, we recommend you to use detectron2 as an library and take
this file as an example of how to use the library.
You may want to write your own script with your datasets and other customizations.
"""
import numpy as np
from PIL import Image 
import PIL 
import cv2
import matplotlib.pyplot as plt
import logging
import os
import time
from collections import OrderedDict
import torch
from detectron2.data.detection_utils import read_image
import detectron2.utils.comm as comm
from detectron2.checkpoint import DetectionCheckpointer
from detectron2.config import get_cfg
from detectron2.utils.visualizer import ColorMode, Visualizer
import detectron2.data.detection_utils as utils
import detectron2.data.transforms as T
from detectron2.data import (
    MetadataCatalog,
)
from detectron2.engine import DefaultTrainer, default_argument_parser, default_setup, hooks, launch
from detectron2.evaluation import (
    CityscapesInstanceEvaluator,
    CityscapesSemSegEvaluator,
    COCOEvaluator,
    COCOPanopticEvaluator,
    DatasetEvaluators,
    LVISEvaluator,
    PascalVOCDetectionEvaluator,
    SemSegEvaluator,
    verify_results,
    NWPUDetectionEvaluator,
    DIORDetectionEvaluator,
    HRRSDDetectionEvaluator,
)
###
import sys
from grad_cam_new import GradCAM
#打印该文件的路径
PROJECT_ABSOLUTE_PATH=os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
# print(PROJECT_ABSOLUTE_PATH)
sys.path.append(os.path.join(PROJECT_ABSOLUTE_PATH))
###
import wsl.data.datasets
from wsl.config import add_wsl_config
from wsl.modeling import GeneralizedRCNNWithTTAAVG, GeneralizedRCNNWithTTAUNION


class Trainer(DefaultTrainer):
    """
    We use the "DefaultTrainer" which contains pre-defined default logic for
    standard training workflow. They may not work for you, especially if you
    are working on a new research project. In that case you can write your
    own training loop. You can use "tools/plain_train_net.py" as an example.
    """

    def __init__(self, cfg):
        super().__init__(cfg)

        self.iter_size = cfg.WSL.ITER_SIZE

        # if comm.is_main_process():
        # self._hooks[-1]._period = cfg.WSL.ITER_SIZE * 10

    def run_step(self):
        """
        Implement the standard training logic described above.
        """
        assert self.model.training, "[SimpleTrainer] model was changed to eval mode!"
        start = time.perf_counter()
        """
        If your want to do something with the data, you can wrap the dataloader.
        """
        #原始代码
        while True:
            data = next(self._data_loader_iter)

            has_gt = [len(x["instances"]) > 0 for x in data]
            if all(has_gt):
                break
            else:
                print(data)

        #去掉了while True
#         data = next(self._data_loader_iter)

        data_time = time.perf_counter() - start

        """
        If your want to do something with the losses, you can wrap the model.
        """
        loss_dict = self.model(data)

#给cls loss 和 refine loss加上权重系数，加大cls loss的权重，这样应该可以增加定位的准确性
#         losses = 0
#         for loss in loss_dict:
# #             print(loss)
#             if loss == 'loss_cls':
#                 losses += 2*loss_dict[loss]
#             else:
#                 losses += loss_dict[loss]
#原始代码，loss中间的权重为1
        losses = sum(loss for loss in loss_dict.values())
        self._detect_anomaly(losses, loss_dict)

        metrics_dict = loss_dict
        metrics_dict["data_time"] = data_time
        self._write_metrics(metrics_dict)

        """
        If you need accumulate gradients or something similar, you can
        wrap the optimizer with your custom `zero_grad()` method.
        """
        if self.iter == self.start_iter:
            self.optimizer.zero_grad()

        losses = losses / self.iter_size
        losses.backward()

        if self.iter % self.iter_size == 0:
            """
            If you need gradient clipping/scaling or other processing, you can
            wrap the optimizer with your custom `step()` method.
            """
            self.optimizer.step()

            self.optimizer.zero_grad()

        del losses
        del loss_dict
        torch.cuda.empty_cache()

    @classmethod
    def build_evaluator(cls, cfg, dataset_name, output_folder=None):
        """
        Create evaluator(s) for a given dataset.
        This uses the special metadata "evaluator_type" associated with each builtin dataset.
        For your own dataset, you can simply create an evaluator manually in your
        script and do not have to worry about the hacky if-else logic here.
        """
        if output_folder is None:
            output_folder = os.path.join(cfg.OUTPUT_DIR, "inference")
        evaluator_list = []
        evaluator_type = MetadataCatalog.get(dataset_name).evaluator_type
        if evaluator_type in ["sem_seg", "coco_panoptic_seg"]:
            evaluator_list.append(
                SemSegEvaluator(
                    dataset_name,
                    distributed=True,
                    num_classes=cfg.MODEL.SEM_SEG_HEAD.NUM_CLASSES,
                    ignore_label=cfg.MODEL.SEM_SEG_HEAD.IGNORE_VALUE,
                    output_dir=output_folder,
                )
            )
        if evaluator_type in ["coco", "coco_panoptic_seg"]:
            evaluator_list.append(COCOEvaluator(dataset_name, cfg, True, output_folder))
        if evaluator_type == "coco_panoptic_seg":
            evaluator_list.append(COCOPanopticEvaluator(dataset_name, output_folder))
        if evaluator_type == "cityscapes_instance":
            assert (
                torch.cuda.device_count() >= comm.get_rank()
            ), "CityscapesEvaluator currently do not work with multiple machines."
            return CityscapesInstanceEvaluator(dataset_name)
        if evaluator_type == "cityscapes_sem_seg":
            assert (
                torch.cuda.device_count() >= comm.get_rank()
            ), "CityscapesEvaluator currently do not work with multiple machines."
            return CityscapesSemSegEvaluator(dataset_name)
        elif evaluator_type == "pascal_voc":
            return PascalVOCDetectionEvaluator(dataset_name)
        elif evaluator_type == "nwpu":
            return NWPUDetectionEvaluator(dataset_name)
        elif evaluator_type == "dior":
            return DIORDetectionEvaluator(dataset_name)  
        elif evaluator_type == "hrrsd":
            return HRRSDDetectionEvaluator(dataset_name) 
        elif evaluator_type == "lvis":
            return LVISEvaluator(dataset_name, cfg, True, output_folder)
        if len(evaluator_list) == 0:
            raise NotImplementedError(
                "no Evaluator for the dataset {} with the type {}".format(
                    dataset_name, evaluator_type
                )
            )
        elif len(evaluator_list) == 1:
            return evaluator_list[0]
        return DatasetEvaluators(evaluator_list)

    @classmethod
    def test_with_TTA(cls, cfg, model):
        if cfg.TEST.EVAL_TRAIN:
            cfg.defrost()
            DATASETS_TEST = cfg.DATASETS.TEST
            DATASETS_PROPOSAL_FILES_TEST = cfg.DATASETS.PROPOSAL_FILES_TEST
            cfg.DATASETS.TEST = cfg.DATASETS.TEST + cfg.DATASETS.TRAIN
            cfg.DATASETS.PROPOSAL_FILES_TEST = (
                cfg.DATASETS.PROPOSAL_FILES_TEST + cfg.DATASETS.PROPOSAL_FILES_TRAIN
            )
            cfg.freeze()

        logger = logging.getLogger("detectron2.trainer")
        # In the end of training, run an evaluation with TTA
        # Only support some R-CNN models.
        logger.info("Running inference with test-time augmentation ...")
        model = GeneralizedRCNNWithTTAAVG(cfg, model)
        evaluators = [
            cls.build_evaluator(
                cfg, name, output_folder=os.path.join(cfg.OUTPUT_DIR, "inference_TTA")
            )
            for name in cfg.DATASETS.TEST
        ]
        res = cls.test(cfg, model, evaluators)
        res = OrderedDict({k + "_TTA": v for k, v in res.items()})

        if cfg.TEST.EVAL_TRAIN:
            cfg.defrost()
            cfg.DATASETS.TEST = DATASETS_TEST
            cfg.DATASETS.PROPOSAL_FILES_TEST = DATASETS_PROPOSAL_FILES_TEST
            cfg.freeze()

        return res
    
    @classmethod
    def vis_box(cls, cfg, model):
        os.makedirs(cfg.VIS_BOX_DIR, exist_ok=True)
        model.eval()
        metadata = MetadataCatalog.get(
            cfg.DATASETS.TEST[0] if len(cfg.DATASETS.TEST) else "__unused"
        )
        for idx, dataset_name in enumerate(cfg.DATASETS.TEST):
            
            data_loader = cls.build_test_loader(cfg, dataset_name)
            for idx, inputs in enumerate(data_loader):
                print(idx)
                for input_ in inputs:
                    img_path = input_['file_name']
                    img = read_image(img_path, format="RGB")
                    predictions  = model(inputs)
#                     print(predictions)
                    visualizer = Visualizer(img, metadata, instance_mode=ColorMode.SEGMENTATION)
                    if "instances" in predictions[0]:
                        print('111111')
                        instances = predictions[0]["instances"].to(torch.device("cpu"))
                        vis_output = visualizer.draw_instance_predictions(predictions=instances)
                        out_filename = os.path.join(cfg.VIS_BOX_DIR, os.path.basename(img_path))
                        print(out_filename)
                        vis_output.save(out_filename)
                        
    @classmethod
    def vis_cam(cls,cfg,model):
        def show_cam_on_image(img, mask, out_dir, idx, category):
            category = category.cpu().numpy()
            
            heatmap = cv2.applyColorMap(np.uint8(255*mask), cv2.COLORMAP_JET)
            heatmap = np.float32(heatmap) 
            
#             img = cv2.imread(img_path)
            #img的读取是用PIL读的RGB格式，现在转化成BGR
            img = img[:,:,::-1]
            #0.5用于削弱heatmap的显示
            cam = 0.7*heatmap + np.float32(img)
            cam = cam / np.max(cam)

            path_cam_img = os.path.join(out_dir, "cam_img_{}_{}.jpg".format(str(idx),str(category)))
            path_raw_img = os.path.join(out_dir, "raw_img_{}_{}.jpg".format(str(idx),str(category)))
            path_raw_cam = os.path.join(out_dir, "raw_cam_{}_{}.jpg".format(str(idx),str(category)))
            
            os.makedirs(out_dir, exist_ok=True)
            
            cv2.imwrite(path_cam_img, np.uint8(255 * cam))
            cv2.imwrite(path_raw_img, np.uint8(img))
            cv2.imwrite(path_raw_cam, np.uint8(heatmap))
            
        layer_name = 'backbone.plain5.0.conv3'
        grad_cam_type = "GradCAM"
        is_train = False
#         train_or_test_data = 'test'
        print(cfg.DATASETS.TEST)
        print(cfg.DATASETS.TRAIN)
#         dataset = cfg.DATASETS.TEST + cfg.DATASETS.TRAIN#cfg.DATASETS.TRAIN or cfg.DATASETS.TEST
#         dataset = [cfg.DATASETS.TEST, cfg.DATASETS.TRAIN]
        if is_train == True:
            dataset = cfg.DATASETS.TRAIN
        else:
            dataset = cfg.DATASETS.TEST
        os.makedirs(cfg.VIS_CAM_DIR, exist_ok=True)
        metadata = MetadataCatalog.get(
            dataset[0] if len(dataset) else "__unused"
        )

#         看每个module的name
        for (name, module) in model.named_modules():
            print(name,module)

        for idx, dataset_name in enumerate(dataset):
#             try:
#                 data_loader = cls.build_train_loader(cfg)
#             except:
#             if dataset_name == cfg.DATASETS.TRAIN:
#                 print('loading train dataset')
#                data_loader = cls.build_train_loader(cfg)   
#             elif dataset_name == cfg.DATASETS.TEST:
#                 print('loading test dataset')
            if is_train == True:
                data_loader = cls.build_train_loader(cfg) 
            else:
                data_loader = cls.build_test_loader(cfg, dataset_name)
#             data_loader = cls.build_test_loader(cfg, dataset_name)
#             print('1111111',len(data_loader))
            for idx, inputs in enumerate(data_loader):
                print(idx)
                for input_ in inputs:
#                     print(input_.keys())
                    img_path = input_['file_name']

                    img_name = os.path.basename(img_path).split('.')[0]
#                     print(img_name)
#                     print(img_path)
                    img = read_image(img_path, format="RGB")

                    #为了让图片和cam的图显示一致，要用到以下操作
#                     cfg.INPUT.MIN_SIZE_TRAIN = (400,)
#                     aug_input = T.StandardAugInput(img)
#                     augs = utils.build_augmentation(cfg, is_train)
#                     aug_input.apply_augmentations(augs)
#                     img = aug_input.image

                    if grad_cam_type == "GradCAM":
                        grad_cam = GradCAM(model, layer_name)
                    elif grad_cam_type == "GradCAM++":
                        grad_cam = GradCamPlusPlus(model, layer_name)
                    else:
                        raise ValueError('Grad CAM type not specified')
#                     predictions  = model(inputs)

#                     print(111111)
                    cam,category = grad_cam(inputs,None)
                    show_cam_on_image(img,cam,cfg.VIS_CAM_DIR,img_name,category)


        

def setup(args):
    """
    Create configs and perform basic setups.
    """
    cfg = get_cfg()
    add_wsl_config(cfg)
    cfg.merge_from_file(args.config_file)
    cfg.merge_from_list(args.opts)
    #freeze()用于冻结cfg里的配置，防止后期被更改。
#     cfg.freeze()
    default_setup(cfg, args)
    return cfg


def main(args):
    cfg = setup(args)

    #测试
    if args.eval_only:
        model = Trainer.build_model(cfg)
        DetectionCheckpointer(model, save_dir=cfg.OUTPUT_DIR).resume_or_load(
            cfg.MODEL.WEIGHTS, resume=args.resume
        )
        res = Trainer.test(cfg, model)
#         res = Trainer.test_with_TTA(cfg, model)
#         if cfg.TEST.AUG.ENABLED:
#             res.update(Trainer.test_with_TTA(cfg, model))
        if comm.is_main_process():
            verify_results(cfg, res)
        return res
    
    #推理画框
    if cfg.VIS_BOX_DIR:
        print('Visualizing boxes')
        model = Trainer.build_model(cfg)
        DetectionCheckpointer(model, save_dir=cfg.VIS_BOX_DIR).resume_or_load(
            cfg.MODEL.WEIGHTS, resume=args.resume
        )
        Trainer.vis_box(cfg, model)
        return

    #画Cam图
    if cfg.VIS_CAM_DIR:
        print('Visualizing cams')
        model = Trainer.build_model(cfg)
        DetectionCheckpointer(model).load(cfg.MODEL.WEIGHTS)
        Trainer.vis_cam(cfg, model)
        return
        
    """
    If you'd like to do anything fancier than the standard training logic,
    consider writing your own training loop (see plain_train_net.py) or
    subclassing the trainer.
    """
    trainer = Trainer(cfg)
    trainer.resume_or_load(resume=args.resume)
    if cfg.TEST.AUG.ENABLED:
        trainer.register_hooks(
            [hooks.EvalHook(0, lambda: trainer.test_with_TTA(cfg, trainer.model))]
        )
    return trainer.train()


if __name__ == "__main__":
    args = default_argument_parser().parse_args()
    print("Command Line Args:", args)
    launch(
        main,
        args.num_gpus,
        num_machines=args.num_machines,
        machine_rank=args.machine_rank,
        dist_url=args.dist_url,
        args=(args,),
    )
